import dotenv
from langchain_core.messages import ToolCall, AIMessage, ToolMessage, HumanMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableConfig
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from typing import Any
dotenv.load_dotenv()
class CustomToolExceptiong(Exception):
    """自定义工具错误异常"""
    def __init__(self,tool_call:ToolCall,exception:Exception)->None:
        super().__init__()
        self.tool_call = tool_call
        self.exception = exception
@tool
def complex_tool(int_arg: int, float_arg: float, dict_arg: dict) -> int:
    """使用复杂工具进行计算操作"""
    return int(int_arg * float_arg)
def tool_custom_exception(msg:AIMessage,config:RunnableConfig)->Any:
    try:
        return  complex_tool.invoke(msg.tool_calls[0]["args"],config)
    except Exception as e:
        raise CustomToolExceptiong(msg.tool_calls[0],e)
def exception_to_message(inputs:dict)->dict:
    #1、从输入中提取错误信息
    exception = inputs.pop("exception")
    # 2、将历史消息添加到原始输入中，以使模型知道它上次一犯什么错
    messages =[
        AIMessage(content="",tool_calls = [exception.tool_call]),
        ToolMessage(tool_call_id=exception.tool_call["id"],content=str(exception.exception)),
        HumanMessage(content="最后一次工具调用引发了异常，调用参数应该是5，2.1，{'a','c'}}")
    ]
    inputs["last_output"] = messages
    return inputs
# 1、创建prompt 并预留占位符，用于存储错误信息
prompt = ChatPromptTemplate.from_messages([
    ("human","{query}"),
    ("placeholder","{last_output}")
])

# 1、创建大模型语言模型并绑定工具
llm = ChatOpenAI(model = "gpt-4o",temperature=0).bind_tools(tools = [complex_tool],tool_choice="any")

# 2、创建链并执行工具
chain = prompt | llm | tool_custom_exception
self_correct_chain = chain.with_fallbacks(
    [exception_to_message|chain],exception_key="exception"
)
# 3、调用链
print(self_correct_chain.invoke({"query":"使用复杂工具，对应参数为5和2.1"}))